from dojo.models import Finding


class DeepfenceThreatmapperVulnerability:
    def get_findings(self, row, headers, test):
        if "cve_id" in headers and "cve_attack_vector" in headers:
            return self._parse_old_format(row, headers, test)
        if "CVE ID" in headers and "Attack Vector" in headers:
            return self._parse_new_format(row, headers, test)
        return None

    def _parse_old_format(self, row, headers, test):
        cve_attack_vector = row[headers["cve_attack_vector"]]
        cve_caused_by_package = row[headers["cve_caused_by_package"]]
        cve_container_image = row[headers["cve_container_image"]]
        cve_container_image_id = row[headers["cve_container_image_id"]]
        cve_description = row[headers["cve_description"]]
        cve_fixed_in = row[headers["cve_fixed_in"]]
        cve_id = row[headers["cve_id"]]
        cve_link = row[headers["cve_link"]]
        cve_severity = row[headers["cve_severity"]]
        cve_overall_score = row[headers["cve_overall_score"]]
        cve_type = row[headers["cve_type"]]
        host_name = row[headers["host_name"]]
        cloud_account_id = row[headers["cloud_account_id"]]
        masked = row[headers["masked"]]

        description = (
            f"**Attack Vector:** {cve_attack_vector}\n"
            f"**Caused By Package:** {cve_caused_by_package}\n"
            f"**Container Image:** {cve_container_image}\n"
            f"**Container Image ID:** {cve_container_image_id}\n"
            f"**Description:** {cve_description}\n"
            f"**Severity:** {cve_severity}\n"
            f"**Overall Score:** {cve_overall_score}\n"
            f"**Type:** {cve_type}\n"
            f"**Host Name:** {host_name}\n"
            f"**Cloud Account ID:** {cloud_account_id}\n"
            f"**Masked:** {masked}\n"
        )

        return Finding(
            title=f"Threatmapper_Vuln_Report-{cve_id}",
            description=description,
            component_name=cve_caused_by_package,
            severity=self.severity(cve_severity),
            static_finding=False,
            dynamic_finding=True,
            mitigation=cve_fixed_in,
            references=cve_link,
            cve=cve_id,
            test=test,
        )

    def _parse_new_format(self, row, headers, test):
        cve_id = row[headers["CVE ID"]]
        severity = row[headers["Severity"]]
        attack_vector = row[headers["Attack Vector"]]
        caused_by_package = row[headers["Caused By Package"]]
        caused_by_package_path = row[headers["Caused By Package Path"]]
        cvss_score = row[headers["CVSS Score"]]
        description_text = row[headers["Description"]]
        fixed_in = row[headers["Fixed In"]]
        link = row[headers["Link"]]
        overall_score = row[headers["Overall Score"]]
        cve_type = row[headers["Type"]]
        node_name = row[headers["Node Name"]]
        node_type = row[headers["Node Type"]]
        cluster_name = row[headers["Kubernetes Cluster Name"]]
        masked = row[headers["Masked"]]

        description = (
            f"**CVE ID:** {cve_id}\n"
            f"**Severity:** {severity}\n"
            f"**Attack Vector:** {attack_vector}\n"
            f"**Caused By Package:** {caused_by_package}\n"
            f"**Caused By Package Path:** {caused_by_package_path}\n"
            f"**CVSS Score:** {cvss_score}\n"
            f"**Description:** {description_text}\n"
            f"**Fixed In:** {fixed_in}\n"
            f"**Link:** {link}\n"
            f"**Overall Score:** {overall_score}\n"
            f"**Type:** {cve_type}\n"
            f"**Node Name:** {node_name}\n"
            f"**Node Type:** {node_type}\n"
            f"**Kubernetes Cluster Name:** {cluster_name}\n"
            f"**Masked:** {masked}\n"
        )

        return Finding(
            title=f"Threatmapper_Vuln_Report-{cve_id}",
            description=description,
            component_name=caused_by_package,
            severity=self.severity(severity),
            static_finding=False,
            dynamic_finding=True,
            mitigation=fixed_in,
            references=link,
            cve=cve_id,
            test=test,
        )

    def severity(self, severity_input):
        if severity_input is None:
            return "Info"
        return severity_input.capitalize()
